package com.keya.features.tool;

import cn.hutool.core.util.ObjectUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.dcm4che3.data.*;
import org.dcm4che3.io.DicomInputStream;
import org.dcm4che3.io.DicomOutputStream;
import org.dcm4che3.io.SAXReader;
import org.dcm4che3.io.SAXWriter;
import org.dcm4che3.net.*;
import org.dcm4che3.net.pdu.AAssociateRQ;
import org.dcm4che3.net.pdu.ExtendedNegotiation;
import org.dcm4che3.net.pdu.PresentationContext;
import org.dcm4che3.tool.common.CLIUtils;
import org.dcm4che3.util.SafeClose;
import org.dcm4che3.util.TagUtils;

import javax.xml.transform.OutputKeys;
import javax.xml.transform.Templates;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.sax.SAXTransformerFactory;
import javax.xml.transform.sax.TransformerHandler;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;
import java.io.*;
import java.security.GeneralSecurityException;
import java.text.DecimalFormat;
import java.text.MessageFormat;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * FindSCU
 *
 * @author : xf.he
 * @version : 1.0 2023/11/2
 * @since : 1.0
 **/
@Data
@Slf4j
public class FindSCU {

    public enum InformationModel {
        PatientRoot(UID.PatientRootQueryRetrieveInformationModelFind, "STUDY"),
        StudyRoot(UID.StudyRootQueryRetrieveInformationModelFind, "STUDY"),
        PatientStudyOnly(UID.PatientStudyOnlyQueryRetrieveInformationModelFind, "STUDY"),
        MWL(UID.ModalityWorklistInformationModelFind, null),
        UPSPull(UID.UnifiedProcedureStepPull, null),
        UPSWatch(UID.UnifiedProcedureStepWatch, null),
        UPSQuery(UID.UnifiedProcedureStepQuery, null),
        HangingProtocol(UID.HangingProtocolInformationModelFind, null),
        ColorPalette(UID.ColorPaletteQueryRetrieveInformationModelFind, null);

        final String cuid;
        final String level;

        InformationModel(String cuid, String level) {
            this.cuid = cuid;
            this.level = level;
        }

        public void adjustQueryOptions(EnumSet<QueryOption> queryOptions) {
            if (level == null) {
                queryOptions.add(QueryOption.RELATIONAL);
                queryOptions.add(QueryOption.DATETIME);
            }
        }
    }

    private static ResourceBundle rb = ResourceBundle.getBundle("dcm4che.findscu.messages");
    private static SAXTransformerFactory saxtf;
    private final Device device = new Device("findscu");
    private final ApplicationEntity ae = new ApplicationEntity("FINDSCU");
    private final Connection conn = new Connection();
    private final Connection remote = new Connection();
    private final AAssociateRQ rq = new AAssociateRQ();
    private int priority;
    private int cancelAfter;
    private InformationModel model;
    private File outDir;
    private DecimalFormat outFileFormat;
    private int[] inFilter;
    private Attributes keys = new Attributes();
    private boolean catOut = false;
    private boolean xml = false;
    private boolean xmlIndent = false;
    private boolean xmlIncludeKeyword = true;
    private boolean xmlIncludeNamespaceDeclaration = false;
    private File xsltFile;
    private Templates xsltTpls;
    private OutputStream out;
    private Association as;
    private AtomicInteger totNumMatches = new AtomicInteger();
    private List<FindResult> findResults = new ArrayList<>();
    private List<FindStudyResult> findStudyResults = new ArrayList<>();
    private Set<String> includeFields;

    public FindSCU() {
        device.addConnection(conn);
        device.addApplicationEntity(ae);
        ae.addConnection(conn);
    }

    public final void setInformationModel(InformationModel model, String[] tss, EnumSet<QueryOption> queryOptions) {
        this.model = model;
        if (rq.getPresentationContext(1) == null) {
            rq.addPresentationContext(new PresentationContext(1, model.cuid, tss));
        }
        if (!queryOptions.isEmpty()) {
            model.adjustQueryOptions(queryOptions);
            rq.addExtendedNegotiation(new ExtendedNegotiation(model.cuid,
                    QueryOption.toExtendedNegotiationInformation(queryOptions)));
        }
        if (model.level != null)
            addLevel(model.level);
    }

    public void addLevel(String s) {
        keys.setString(Tag.QueryRetrieveLevel, VR.CS, s);
    }


    public final void setOutputDirectory(File outDir) {
        outDir.mkdirs();
        this.outDir = outDir;
    }

    private static CommandLine parseComandLine(String[] args) throws ParseException {
        Options opts = new Options();
        addServiceClassOptions(opts);
        addKeyOptions(opts);
        addOutputOptions(opts);
        addQueryLevelOption(opts);
        addCancelOption(opts);
        CLIUtils.addConnectOption(opts);
        CLIUtils.addBindOption(opts, "FINDSCU");
        CLIUtils.addAEOptions(opts);
        CLIUtils.addSendTimeoutOption(opts);
        CLIUtils.addResponseTimeoutOption(opts);
        CLIUtils.addPriorityOption(opts);
        CLIUtils.addCommonOptions(opts);
        return CLIUtils.parseComandLine(args, opts, rb, FindSCU.class);
    }

    private static void addServiceClassOptions(Options opts) {
        opts.addOption(Option.builder("M")
                .hasArg()
                .argName("name")
                .desc(rb.getString("model"))
                .build());
        CLIUtils.addTransferSyntaxOptions(opts);
        opts.addOption(null, "relational", false, rb.getString("relational"));
        opts.addOption(null, "datetime", false, rb.getString("datetime"));
        opts.addOption(null, "fuzzy", false, rb.getString("fuzzy"));
        opts.addOption(null, "timezone", false, rb.getString("timezone"));
    }

    private static void addQueryLevelOption(Options opts) {
        opts.addOption(Option.builder("L")
                .hasArg()
                .argName("PATIENT|STUDY|SERIES|IMAGE")
                .desc(rb.getString("level"))
                .build());
    }

    private static void addCancelOption(Options opts) {
        opts.addOption(Option.builder()
                .longOpt("cancel")
                .hasArg()
                .argName("num-matches")
                .desc(rb.getString("cancel"))
                .build());
    }

    private static void addKeyOptions(Options opts) {
        opts.addOption(Option.builder("m")
                .hasArgs()
                .argName("[seq.]attr=value")
                .desc(rb.getString("match"))
                .build());
        opts.addOption(Option.builder("r")
                .hasArgs()
                .argName("[seq.]attr")
                .desc(rb.getString("return"))
                .build());
        opts.addOption(Option.builder("i")
                .hasArgs()
                .argName("attr")
                .desc(rb.getString("in-attr"))
                .build());
    }

    private static void addOutputOptions(Options opts) {
        opts.addOption(Option.builder()
                .longOpt("out-dir")
                .hasArg()
                .argName("directory")
                .desc(rb.getString("out-dir"))
                .build());
        opts.addOption(Option.builder()
                .longOpt("out-file")
                .hasArg()
                .argName("name")
                .desc(rb.getString("out-file"))
                .build());
        opts.addOption("X", "xml", false, rb.getString("xml"));
        opts.addOption(Option.builder("x")
                .longOpt("xsl")
                .hasArg()
                .argName("xsl-file")
                .desc(rb.getString("xsl"))
                .build());
        opts.addOption("I", "indent", false, rb.getString("indent"));
        opts.addOption("K", "no-keyword", false, rb.getString("no-keyword"));
        opts.addOption(null, "xmlns", false, rb.getString("xmlns"));
        opts.addOption(null, "out-cat", false, rb.getString("out-cat"));
    }

    private static EnumSet<QueryOption> queryOptionsOf(CommandLine cl) {
        EnumSet<QueryOption> queryOptions = EnumSet.noneOf(QueryOption.class);
        if (cl.hasOption("relational"))
            queryOptions.add(QueryOption.RELATIONAL);
        if (cl.hasOption("datetime"))
            queryOptions.add(QueryOption.DATETIME);
        if (cl.hasOption("fuzzy"))
            queryOptions.add(QueryOption.FUZZY);
        if (cl.hasOption("timezone"))
            queryOptions.add(QueryOption.TIMEZONE);
        return queryOptions;
    }

    private void configureOutput(CommandLine cl) {
        if (cl.hasOption("out-dir"))
            setOutputDirectory(new File(cl.getOptionValue("out-dir")));
        setOutFileFormat(new DecimalFormat(cl.getOptionValue("out-file", "000'.dcm'")));
        setCatOut(cl.hasOption("out-cat"));
        setXml(cl.hasOption("X"));
        if (cl.hasOption("x")) {
            setXml(true);
            setXsltFile(new File(cl.getOptionValue("x")));
        }
        setXmlIndent(cl.hasOption("I"));
        setXmlIncludeKeyword(!cl.hasOption("K"));
        setXmlIncludeNamespaceDeclaration(cl.hasOption("xmlns"));
    }

    private void configureCancel(CommandLine cl) {
        if (cl.hasOption("cancel"))
            setCancelAfter(Integer.parseInt(cl.getOptionValue("cancel")));
    }

    private void configureKeys(CommandLine cl) {
        CLIUtils.addEmptyAttributes(keys, cl.getOptionValues("r"));
        CLIUtils.addAttributes(keys, cl.getOptionValues("m"));
        if (cl.hasOption("L"))
            addLevel(cl.getOptionValue("L"));
        if (cl.hasOption("i"))
            setInFilter(CLIUtils.toTags(cl.getOptionValues("i")));
    }

    private void configureServiceClass(CommandLine cl) throws ParseException {
        setInformationModel(informationModelOf(cl),
                CLIUtils.transferSyntaxesOf(cl), queryOptionsOf(cl));
    }

    private static InformationModel informationModelOf(CommandLine cl) throws ParseException {
        try {
            return cl.hasOption("M")
                    ? InformationModel.valueOf(cl.getOptionValue("M"))
                    : InformationModel.StudyRoot;
        } catch (IllegalArgumentException e) {
            throw new ParseException(
                    MessageFormat.format(
                            rb.getString("invalid-model-name"),
                            cl.getOptionValue("M")));
        }
    }

    public void open() throws IOException, InterruptedException, IncompatibleConnectionException, GeneralSecurityException {
        as = ae.connect(conn, remote, rq);
    }

    public void close() throws IOException, InterruptedException {
        if (as != null && as.isReadyForDataTransfer()) {
            as.waitForOutstandingRSP();
            as.release();
        }
        SafeClose.close(out);
        out = null;
    }

    public void query(File f) throws Exception {
        Attributes attrs;
        String filePath = f.getPath();
        String fileExt = filePath.substring(filePath.lastIndexOf(".") + 1).toLowerCase();
        DicomInputStream dis = null;
        try {
            attrs = fileExt.equals("xml")
                    ? SAXReader.parse(filePath)
                    : new DicomInputStream(f).readDataset();
            if (inFilter != null) {
                attrs = new Attributes(inFilter.length + 1);
                attrs.addSelected(attrs, inFilter);
            }
        } finally {
            SafeClose.close(dis);
        }
        mergeKeys(attrs, keys);
        query(attrs);
    }

    private static class MergeNested implements Attributes.Visitor {
        private final Attributes keys;

        MergeNested(Attributes keys) {
            this.keys = keys;
        }

        @Override
        public boolean visit(Attributes attrs, int tag, VR vr, Object val) {
            if (isNotEmptySequence(val)) {
                Object o = keys.remove(tag);
                if (isNotEmptySequence(o))
                    ((Sequence) val).get(0).addAll(((Sequence) o).get(0));
            }
            return true;
        }

        private static boolean isNotEmptySequence(Object val) {
            return val instanceof Sequence && !((Sequence) val).isEmpty();
        }
    }

    static void mergeKeys(Attributes attrs, Attributes keys) {
        try {
            attrs.accept(new MergeNested(keys), false);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        attrs.addAll(keys);
    }

    public void query() throws IOException, InterruptedException {
        query(keys);
    }

    private void query(Attributes keys) throws IOException, InterruptedException {
        DimseRSPHandler rspHandler = new DimseRSPHandler(as.nextMessageID()) {

            int cancelAfter = FindSCU.this.cancelAfter;
            int numMatches;

            @Override
            public void onDimseRSP(Association as, Attributes cmd,
                                   Attributes data) {
                super.onDimseRSP(as, cmd, data);
                int status = cmd.getInt(Tag.Status, -1);
                if (Status.isPending(status)) {
                    FindSCU.this.onResult(data);
                    ++numMatches;
                    if (cancelAfter != 0 && numMatches >= cancelAfter) {
                        try {
                            cancel(as);
                            cancelAfter = 0;
                        } catch (IOException e) {
                            log.error("DIMSERSP-" + "cancel" + "error", e);
                        }
                    }
                    try {
                        setFindResult(data);
                        setFindResult(data, includeFields);
                    } catch (Exception e) {
                        log.error("setFindResult exception", e);
                    }
                }
            }
        };
        query(keys, rspHandler);
    }

    private void setFindResult(Attributes attrs) {
        FindStudyResult result = new FindStudyResult();
        result.setAccessionNumber(attrs.getString(Tag.AccessionNumber))
                .setPatientName(attrs.getString(Tag.PatientName))
                .setPatientID(attrs.getString(Tag.PatientID))
                .setPatientSex(attrs.getString(Tag.PatientSex))
                .setPatientBirthDate(attrs.getString(Tag.PatientBirthDate))
                .setRetrieveAETitle(attrs.getString(Tag.RetrieveAETitle))
                .setSpecificCharacterSet(attrs.getString(Tag.SpecificCharacterSet))
                .setStudyDate(attrs.getString(Tag.StudyDate))
                .setStudyInstanceUID(attrs.getString(Tag.StudyInstanceUID))
                .setSeriesInstanceUID(attrs.getString(Tag.SeriesInstanceUID))
                .setSeriesNumber(attrs.getString(Tag.SeriesNumber))
                .setSeriesDescription(attrs.getString(Tag.SeriesDescription))
                .setSliceThickness(attrs.getString(Tag.SliceThickness))
                .setBodyPartExamined(attrs.getString(Tag.BodyPartExamined))
                .setNumberOfSlices(attrs.getString(Tag.NumberOfSlices));
        findStudyResults.add(result);
    }

    private void setFindResult(Attributes attrs, Set<String> includeFields) {
        if (ObjectUtil.isEmpty(includeFields)) {
            return;
        }
        Map<String, Object> resultMap = new HashMap<>();
        for (String name : includeFields) {
            int tag = TagUtils.forName(name);
            if (tag == -1) {
                log.error("没获取到这个值:" + name);
                continue;
            }
            Object value;
            VR vr = attrs.getVR(tag);
            if (vr != null && vr.isIntType()) {
                value = attrs.getInt(tag, -1);
            } else {
                value = attrs.getString(tag);
            }
            resultMap.put(name, value);
        }
        FindResult findResult = new FindResult(resultMap);
        findResults.add(findResult);
    }


    public void query(DimseRSPHandler rspHandler) throws IOException, InterruptedException {
        query(keys, rspHandler);
    }

    private void query(Attributes keys, DimseRSPHandler rspHandler) throws IOException, InterruptedException {
        as.cfind(model.cuid, priority, keys, null, rspHandler);
    }

    private void onResult(Attributes data) {
        int numMatches = totNumMatches.incrementAndGet();
        if (outDir == null)
            return;

        try {
            if (out == null) {
                File f = new File(outDir, fname(numMatches));
                out = new BufferedOutputStream(
                        new FileOutputStream(f));
            }
            if (xml) {
                writeAsXML(data, out);
            } else {
                DicomOutputStream dos = new DicomOutputStream(out, UID.ImplicitVRLittleEndian);
                dos.writeDataset(null, data);
            }
            out.flush();
        } catch (Exception e) {
            e.printStackTrace();
            SafeClose.close(out);
            out = null;
        } finally {
            if (!catOut) {
                SafeClose.close(out);
                out = null;
            }
        }
    }

    private String fname(int i) {
        synchronized (outFileFormat) {
            return outFileFormat.format(i);
        }
    }

    private void writeAsXML(Attributes attrs, OutputStream out) throws Exception {
        TransformerHandler th = getTransformerHandler();
        th.getTransformer().setOutputProperty(OutputKeys.INDENT,
                xmlIndent ? "yes" : "no");
        th.setResult(new StreamResult(out));
        SAXWriter saxWriter = new SAXWriter(th);
        saxWriter.setIncludeKeyword(xmlIncludeKeyword);
        saxWriter.setIncludeNamespaceDeclaration(xmlIncludeNamespaceDeclaration);
        saxWriter.write(attrs);
    }

    private TransformerHandler getTransformerHandler() throws Exception {
        SAXTransformerFactory tf = saxtf;
        if (tf == null)
            saxtf = tf = (SAXTransformerFactory) TransformerFactory
                    .newInstance();
        if (xsltFile == null)
            return tf.newTransformerHandler();

        Templates tpls = xsltTpls;
        if (tpls == null) ;
        xsltTpls = tpls = tf.newTemplates(new StreamSource(xsltFile));

        return tf.newTransformerHandler(tpls);
    }


    private static String[] initArgs(DicomNetParams netParas, Map<String, String> cFindTags) {
        List<String> list = new ArrayList<>();
        list.add("-c");
        list.add(netParas.getCalledAE() + "@" + netParas.getCalledIP() + ":" + netParas.getCalledPort());
        list.add("-L");
        list.add("STUDY");
//        list.add("SERIES");
        list.add("--relational");
        for (Map.Entry<String, String> entry : cFindTags.entrySet()) {
            if (ObjectUtil.isEmpty(entry.getValue())) {
                list.add("-r");
                list.add(entry.getKey());
            } else {
                list.add("-m");
                list.add(entry.getKey() + "=" + entry.getValue());
            }
        }
        return list.toArray(new String[0]);
    }

    private static String[] initArgs(DicomNetParams netParas, Map<String, String> cFindTags, String l) {
        List<String> list = new ArrayList<>();
        list.add("-c");
        list.add(netParas.getCalledAE() + "@" + netParas.getCalledIP() + ":" + netParas.getCalledPort());
        list.add("-L");
        list.add(l);
        list.add("--relational");
        for (Map.Entry<String, String> entry : cFindTags.entrySet()) {
            if (ObjectUtil.isEmpty(entry.getValue())) {
                list.add("-r");
                list.add(entry.getKey());
            } else {
                list.add("-m");
                list.add(entry.getKey() + "=" + entry.getValue());
            }
        }
        return list.toArray(new String[0]);
    }

    public Set<String> buildDefaultIncludeFields() {
        Set<String> set = new HashSet<>();
        set.add(Keyword.valueOf(Tag.StudyInstanceUID));
        set.add(Keyword.valueOf(Tag.StudyDate));
        set.add(Keyword.valueOf(Tag.StudyTime));
        set.add(Keyword.valueOf(Tag.AccessionNumber));
        set.add(Keyword.valueOf(Tag.StudyDescription));
        set.add(Keyword.valueOf(Tag.NumberOfStudyRelatedSeries));
        set.add(Keyword.valueOf(Tag.NumberOfStudyRelatedInstances));
        set.add(Keyword.valueOf(Tag.PatientName));
        set.add(Keyword.valueOf(Tag.PatientID));
        //Series信息
        set.add(Keyword.valueOf(Tag.SeriesInstanceUID));
        set.add(Keyword.valueOf(Tag.SeriesDescription));
        set.add(Keyword.valueOf(Tag.SeriesNumber));
        set.add(Keyword.valueOf(Tag.SliceThickness));
        set.add(Keyword.valueOf(Tag.BodyPartExamined));
        return set;
    }

    public List<FindResult> cFind(DicomNetParams dicomNetParams, Map<String, String> findParams) {
        List<FindResult> findResults;
        try {
            findResults = cFind(dicomNetParams, findParams, buildDefaultIncludeFields());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return findResults;
    }

    public List<FindStudyResult> cFindStudy(DicomNetParams dicomNetParams, Map<String, String> findParams, String l) throws Exception {
        CommandLine cl = parseComandLine(initArgs(dicomNetParams, findParams, l));
        CLIUtils.configureConnect(remote, rq, cl);
        CLIUtils.configureBind(conn, ae, cl);
        CLIUtils.configure(conn, cl);
        remote.setTlsProtocols(conn.getTlsProtocols());
        remote.setTlsCipherSuites(conn.getTlsCipherSuites());
        configureServiceClass(cl);
        configureKeys(cl);
        configureOutput(cl);
        configureCancel(cl);
        setPriority(CLIUtils.priorityOf(cl));
        device.setExecutor(PacsExecutor.scheduledExecutorService);
        device.setScheduledExecutor(PacsExecutor.scheduledExecutorService);
        try {
            open();
            List<String> argList = cl.getArgList();
            if (argList.isEmpty()) {
                query();
            } else {
                for (String arg : argList) {
                    query(new File(arg));
                }
            }
        } finally {
            try {
                close();
            } catch (Exception e) {
                log.error("as close exception", e);
            }
        }
        return findStudyResults;
    }


    public List<FindResult> cFind(DicomNetParams dicomNetParams, Map<String, String> findParams, Set<String> includeFields) throws Exception {
        this.includeFields = includeFields;
        device.setDeviceName(dicomNetParams.getCallingAE());
        ae.setAETitle(dicomNetParams.getCallingAE());
        CommandLine cl = parseComandLine(initArgs(dicomNetParams, findParams));
        CLIUtils.configureConnect(remote, rq, cl);
        CLIUtils.configureBind(conn, ae, cl);
        CLIUtils.configure(conn, cl);
        remote.setTlsProtocols(conn.getTlsProtocols());
        remote.setTlsCipherSuites(conn.getTlsCipherSuites());
        configureServiceClass(cl);
        configureKeys(cl);
        configureOutput(cl);
        configureCancel(cl);
        setPriority(CLIUtils.priorityOf(cl));
        device.setExecutor(PacsExecutor.scheduledExecutorService);
        device.setScheduledExecutor(PacsExecutor.scheduledExecutorService);
        try {
            open();
            List<String> argList = cl.getArgList();
            if (argList.isEmpty()) {
                query();
            } else {
                for (String arg : argList) {
                    query(new File(arg));
                }
            }
        } finally {
            try {
                close();
            } catch (Exception e) {
                log.error("as close exception", e);
            }
        }
        return findResults;
    }
}
